import torch
from transformers import AutoTokenizer
from model.bert_classifier import BertClassifier
from preprocess.dataset import get_dataset, DataType
from configuration import config
from runner.predict import predict

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(config.PRETRAINED_DIR / 'bert-base-chinese')

dataset = get_dataset(type=DataType.TRAIN)
label_feature = dataset.features['label']

model = BertClassifier(freeze_bert=True).to(device)
model.load_state_dict(torch.load(config.MODELS_DIR / 'model.pt'))


def predict_service(text):
    return predict(text, model, tokenizer, label_feature, device)
